#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2012 Tobias Marschall
# 
# This file is part of CLEVER.
# 
# CLEVER is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# CLEVER is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function, division
from optparse import OptionParser
import sys
import os
import shutil
import time
import subprocess

__author__ = "Tobias Marschall"

usage = """%prog [options] <bam-file> <ref.fasta(.gz)> <result-directory>

This tool runs the whole workflow necessary to use CLEVER.

<bam-file>         Input BAM file. All alignments for the same read (pair) must be in
                   subsequent lines. It is highly recommended to allows multiple 
                   alignments per read to avoid spurious predictions. 
<ref.fasta(.gz)>   The reference genome in (gzipped) FASTA format. This is needed to
                   recompute alignment scores (AS tags). If your BAM file does have AS tags
                   such that 10^(AS/-10.0) can be interpreted as the probability of this
                   alignment being correct, use option -a to omit this step.
<result-directory> Directory to be created to store results in. If it already exists, abort
                   unless option -f is given."""

def locate_executeable(exe_dict, name):
	def isexecutable(f):
		return os.path.isfile(f) and os.access(f, os.X_OK)
	for path in os.environ["PATH"].split(os.pathsep):
		f = os.path.join(path, name)
		if isexecutable(f):
			exe_dict[name] = f
			print('Found executable', f, file=sys.stderr)
			return True
	scriptpath = os.path.dirname(os.path.abspath(__file__))
	f = os.path.join(scriptpath, name)
	if isexecutable(f):
		exe_dict[name] = f
		print('Found executable', f, file=sys.stderr)
		return True
	f = os.path.abspath(os.path.join(scriptpath, '..', 'src', name))
	if isexecutable(f):
		exe_dict[name] = f
		print('Found executable', f, file=sys.stderr)
		return True
	print('Could not locate executable \"%s\". It\'s not in your PATH.'%name, file=sys.stderr)
	return False

def execute(commandline):
	exitcode = os.system(commandline)
	if exitcode == 0:
		return True
	else:
		print("Error: non-zero exit code (%d) returned by commandline:"%exitcode, file=sys.stderr)
		print("\"%s\""%commandline, file=sys.stderr)
		return False


makefile_template = """.DELETE_ON_ERROR:

INPUT := $(shell find -name '*.aln-priors.gz')
RESULTS := $(patsubst ./priors.%.aln-priors.gz,clever.%.out,$(INPUT))

all: $(RESULTS)
	echo "DONE."

clever.%.out: priors.%.sorted-priors.gz lengths.mean-and-sd
	echo Starting to run CLEVER for chromosome $*
	zcat $< | {} lengths.mean-and-sd > $@ 2> clever.$*.log 
	echo Finished running CLEVER for chromosome $*

priors.%.sorted-priors.gz: priors.%.aln-priors.gz
	echo Starting to sort priors for chromosome $*
	zcat $< | sort -g -k6 | gzip > $@
	echo Finished sorting priors for chromosome $*
"""

def main():
	parser = OptionParser(usage=usage)
	parser.add_option("-T", action="store", dest="threads", default=1, type=int,
					help="Number of threads to use (default=1).")
	parser.add_option("-f", action="store_true", dest="force", default=False,
					help="Delete old result and working directory first (if present).")
	parser.add_option("-w", action="store", dest="work_dir", default=None,
					help="Working directory (default: <result-directory>/work).")
	parser.add_option("-a", action="store_true", dest="as_tags", default=False,
					help="Do not (re-)compute AS tags. If given, the argument <ref.fasta(.gz)> is ignored.")
	parser.add_option("-B", action="store_true", dest="bwa", default=False,
					help="Indicates that input BAM file was produced by BWA. If given, the script xa2multi.pl (that comes with BWA) is applied (and thus must be in your PATH). Also, reads for which tags X0 and X1 are such that X0+X1>1, but don't have a XA tag, are ignored. We recommend using options \"-n 25 -N 25\" in the \"sampe\" step to allow up to 25 mappings per read.")
	parser.add_option("-I", action="store_true", dest="ignore_x_tags", default=False,
					help="Ignore inconsistencies in X0 and X1 tags.")
	parser.add_option("-k", action="store_true", dest="keep_workdir", default=False,
					help="Keep working directory (default: delete directory when finished).")
	(options, args) = parser.parse_args()
	if (len(args)!=3):
		parser.print_help()
		return 1
	bam_filename = os.path.abspath(args[0])
	ref_filename = os.path.abspath(args[1])
	result_dir = os.path.abspath(args[2])
	
	# find needed executables
	exe_dict = dict()
	print('===== Checking dependencies =====', file=sys.stderr)
	if not os.path.isfile(bam_filename):
		print("Error: file \"%s\" not found."%bam_filename, file=sys.stderr)
		return 1
	if not os.path.isfile(ref_filename):
		print("Error: file \"%s\" not found."%ref_filename, file=sys.stderr)
		return 1
	if not locate_executeable(exe_dict, 'clever'): return 1
	if not locate_executeable(exe_dict, 'add-score-tags-to-bam'): return 1
	if not locate_executeable(exe_dict, 'bam-to-alignment-priors'): return 1
	if not locate_executeable(exe_dict, 'insert-length-histogram'): return 1
	if not locate_executeable(exe_dict, 'split-priors-by-chromosome'): return 1
	if not locate_executeable(exe_dict, 'make'): return 1
	if options.bwa:
		if not locate_executeable(exe_dict, 'xa2multi.pl'): return 1
		if not locate_executeable(exe_dict, 'samtools'): return 1
	print('===== Checking directories =====', file=sys.stderr)
	if os.path.isdir(result_dir):
		if options.force:
			try:
				shutil.rmtree(result_dir)
			except OSError as e:
				print("Error deleting previous result directory:", e, file=sys.stderr)
				return 1
		else:
			print("Error: directory \"%s\" already exists. Move it out of the way or use option -f."%result_dir, file=sys.stderr)
			return 1
	try:
		os.makedirs(result_dir)
	except OSError as e:
		print("Error creating result directory:", e, file=sys.stderr)
		return 1
	if options.work_dir != None:
		work_dir = os.path.abspath(options.work_dir)
		if os.path.exists(work_dir):
			if options.force:
				shutil.rmtree(work_dir)
			else:
				print("Error: directory \"%s\" already exists. Move it out of the way or use option -f."%work_dir, file=sys.stderr)
				return 1
	else:
		work_dir = os.path.join(result_dir, 'work')
	os.makedirs(work_dir)
	print('Result directory:', result_dir, file=sys.stderr)
	print('Working directory:', work_dir, file=sys.stderr)
	os.chdir(work_dir)
	if not options.as_tags:
		print('===== (Re-)computing alignment scores (AS tags) for %s ====='%bam_filename, file=sys.stderr)
		if options.bwa:
			try:
				samtools_view = subprocess.Popen([exe_dict['samtools'], 'view', '-h', bam_filename], stdout=subprocess.PIPE)
				xa2multi = subprocess.Popen([exe_dict['xa2multi.pl']], stdin=samtools_view.stdout, stdout=subprocess.PIPE)
				samtools_unview = subprocess.Popen([exe_dict['samtools'], 'view', '-S', '-b', '-'], stdin=xa2multi.stdout, stdout=subprocess.PIPE)
				add_score_tags = subprocess.Popen([exe_dict['add-score-tags-to-bam'], '-s', ref_filename], stdin=samtools_unview.stdout, stdout=open('input.bam', 'w'))
				if samtools_view.wait() != 0:
					print('Error executing "samtools view".', file=sys.stderr)
					return 1
				if xa2multi.wait() != 0:
					print('Error executing xa2multi.pl.', file=sys.stderr)
					return 1
				if samtools_unview.wait() != 0:
					print('Error executing "samtools view -S -b".', file=sys.stderr)
					return 1
				if add_score_tags.wait() != 0:
					print('Error executing add-score-tags-to-bam.', file=sys.stderr)
					return 1
			except:
				add_score_tags.terminate()
				samtools_unview.terminate()
				xa2multi.terminate()
				samtools_view.terminate()
				raise
		else:
			try:
				add_score_tags = subprocess.Popen([exe_dict['add-score-tags-to-bam'], '-s', ref_filename], stdin=open(bam_filename), stdout=open('input.bam', 'w'))
				if add_score_tags.wait() != 0:
					print('Error executing add-score-tags-to-bam.', file=sys.stderr)
					return 1
			except:
				add_score_tags.terminate()
				raise
		bam_filename = os.path.abspath(os.path.join(work_dir, 'input.bam'))
	else:
		if not execute('ln -s %s input.bam'%bam_filename): return 1
	print('===== Determining insert size distribution =====', file=sys.stderr)
	if not execute('%s -m lengths.mean-and-sd < input.bam > lengths.histogram'%exe_dict['insert-length-histogram']): return 1
	print('===== Generating alignment priors =====', file=sys.stderr)
	bam_to_priors_call = [exe_dict['bam-to-alignment-priors']]
	if options.bwa: bam_to_priors_call.append('-X')
	if options.ignore_x_tags: bam_to_priors_call.append('-I')
	bam_to_priors_call.append('lengths.histogram')
	try:
		bam_to_priors = subprocess.Popen(bam_to_priors_call, stdin=open('input.bam'), stdout=subprocess.PIPE)
		split_priors = subprocess.Popen(['split-priors-by-chromosome', '-z', 'priors'], stdin=bam_to_priors.stdout)
		if bam_to_priors.wait() != 0:
			print('Error executing bam-to-alignment-priors.', file=sys.stderr)
			return 1
		if split_priors.wait() != 0:
			print('Error executing split-priors-by-chromosome.', file=sys.stderr)
			return 1
	except:
		split_priors.terminate()
		bam_to_priors.terminate()
		raise
	print('Done.', file=sys.stderr)
	time.sleep(1)
	print('===== Sort priors for each chromosome and run CLEVER =====', file=sys.stderr)
	makefile = open('Makefile', 'w')
	makefile.write(makefile_template.format(exe_dict['clever']))
	makefile.close()
	if not execute('%s --quiet -j %d'%(exe_dict['make'],options.threads)): return 1
	print('===== Aggregating results =====', file=sys.stderr)
	execute('cat %s > %s'%(os.path.join(work_dir,'*.out'),os.path.join(result_dir,'predictions.txt')))
	log_dir = os.path.join(result_dir,'logs')
	os.makedirs(log_dir)
	execute('cp %s %s'%(os.path.join(work_dir,'*.log'),log_dir))
	if not options.keep_workdir:
		shutil.rmtree(work_dir)
	print('Done.', file=sys.stderr)


if __name__ == '__main__':
	sys.exit(main())
